package cn.iocoder.yudao.module.ai.config;

import jakarta.annotation.Resource;
import org.springframework.ai.model.tool.ToolCallingManager;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Primary;
import org.springframework.http.client.SimpleClientHttpRequestFactory;
import org.springframework.stereotype.Component;
import org.springframework.web.client.RestClient;
import org.springframework.web.reactive.function.client.WebClient;

import java.time.Duration;

@Component
public class ChatModelConfig {


    private String baseUrl = "http://192.168.1.141:11434";
    private String model = "gpt-oss:20b";

    @Resource
    private WebClient.Builder webClientBuilder;


    @Bean
    @Primary
    public OllamaChatModel ollamaChatModel() {


        //大模型使用call方法时超时问题
        SimpleClientHttpRequestFactory simpleClientHttpRequestFactory = new SimpleClientHttpRequestFactory();
        simpleClientHttpRequestFactory.setConnectTimeout(Duration.ofSeconds(1000));
        simpleClientHttpRequestFactory.setReadTimeout(Duration.ofSeconds(1000));
        RestClient.Builder builder = RestClient.builder()
                .requestFactory(simpleClientHttpRequestFactory);
        OllamaApi ollamaApi = OllamaApi.builder()
                .restClientBuilder( builder)
                .webClientBuilder(webClientBuilder)
                .baseUrl(baseUrl)
                .build();

        return OllamaChatModel
                .builder()
                .defaultOptions(OllamaOptions.builder()
                        .model(model)
                        .temperature(0.1)
                        .topP(0.8)
                        .topK(40)
                        .repeatPenalty(1.1)
                        .numPredict(2048)  // 减少预测token数量，提高响应速度
                        .numCtx(4096)      // 减少上下文长度，节省显存
                        .numThread(12)      // 使用8个线程
                        .build())
                .ollamaApi(ollamaApi)
                .toolCallingManager(ToolCallingManager.builder().build())
                .build();
    }

}
